﻿#include <string>
#include <fstream>
#include <vector>
#include <sstream>
#include <iostream>
#include <chrono>
#include <map>
#include <numeric>
#include <unordered_map>
#include <Eigen/Dense>
#include <boost/filesystem.hpp>
#include <opencv2/opencv.hpp>
#include "common/log.h"
#include <exiv2/exiv2.hpp>

namespace fs = boost::filesystem;

static inline std::string padZeros(int number, int width) {
  std::ostringstream oss;
  oss << std::setw(width) << std::setfill('0') << number;
  return oss.str();
}

namespace {

class Exiv2FrameReader {
 public:
  struct FrameTags {
    FrameTags()
    : timestamp("Xmp.exif.REALM.Timestamp"),
      camera_id("Exif.Image.Model"),
      heading("Xmp.exif.REALM.Heading"),
      latitude("Exif.GPSInfo.GPSLatitude"),
      latituderef("Exif.GPSInfo.GPSLatitudeRef"),
      longitude("Exif.GPSInfo.GPSLongitude"),
      longituderef("Exif.GPSInfo.GPSLongitudeRef"),
      altitude("Exif.GPSInfo.GPSAltitude")
    {
    }

    FrameTags(const std::string &timestamp,
              const std::string &camera_id,
              const std::string &heading,
              const std::string &latitude,
              const std::string &latituderef,
              const std::string &longitude,
              const std::string &longituderef,
              const std::string &altitude)
    : timestamp(timestamp),
      camera_id(camera_id),
      heading(heading),
      latitude(latitude),
      latituderef(latituderef),
      longitude(longitude),
      longituderef(longituderef),
      altitude(altitude)
    {
    }

    std::string timestamp;
    std::string camera_id;
    std::string heading;
    std::string latitude;
    std::string latituderef;
    std::string longitude;
    std::string longituderef;
    std::string altitude;
  };

public:
  Exiv2FrameReader() = default;

  void AddExivInfoToImage(const std::string& imagePath, double latitude, double longitude, double altitude) {
    try {
      Exiv2::Image::UniquePtr image = Exiv2::ImageFactory::open(imagePath);
      image->readMetadata();

      Exiv2::ExifData& exifData = image->exifData();

      if (exifData.empty()) {
        exifData = Exiv2::ExifData();
      }

      // 将经度值转换为DMS格式
      int degrees = static_cast<int>(latitude);
      double minutes = (latitude - degrees) * 60;
      int intMinutes = static_cast<int>(minutes);
      double seconds = (minutes - intMinutes) * 60;
      std::string latitudeStr = std::to_string(degrees) + "/1 " +
                                std::to_string(intMinutes) + "/1 " +
                                std::to_string(static_cast<int>(seconds * 1e7)) + "/10000000";

      // 将纬度值转换为DMS格式
      degrees = static_cast<int>(longitude);
      minutes = (longitude - degrees) * 60;
      intMinutes = static_cast<int>(minutes);
      seconds = (minutes - intMinutes) * 60;
      std::string longitudeStr = std::to_string(degrees) + "/1 " +
                                 std::to_string(intMinutes) + "/1 " +
                                 std::to_string(static_cast<int>(seconds * 1e7)) + "/10000000";

      // 设置经纬度信息
      exifData[m_frame_tags_.latitude/*"Exif.GPSInfo.GPSLatitude"*/] = latitudeStr;
      exifData[m_frame_tags_.longitude/*"Exif.GPSInfo.GPSLongitude"*/] = longitudeStr;
      exifData[m_frame_tags_.altitude/*"Exif.GPSInfo.GPSAltitude"*/] =  Exiv2::Rational(static_cast<int>(altitude * 100), 100);

      std::string longitudeRef = (longitude >= 0) ? "E" : "W";
      std::string latitudeRef = (latitude >= 0) ? "N" : "S";
      exifData[m_frame_tags_.longituderef/*"Exif.GPSInfo.GPSLongitudeRef"*/] = longitudeRef;
      exifData[m_frame_tags_.latituderef/*"Exif.GPSInfo.GPSLatitudeRef"*/] = latitudeRef;

      // 写入更新的EXIF信息
      image->setExifData(exifData);
      image->writeMetadata();
      ADEBUG("Processed image : {}", imagePath);
    } catch (Exiv2::Error& e) {
      AERROR("Error: {}", e.what());
    }
  }

private:
  //! Tags to be read from the exiv information
  FrameTags m_frame_tags_;
};

struct FrameInfo {
  unsigned int rel_timestamp_;  // relative timestamp in ms
  unsigned int diff_time_;  // ms
  std::chrono::milliseconds milli_tm_;
  double focal_length_;
  double latitude_;
  double longitude_;
  double rel_altitude_, abs_alt_;
  double gb_yaw_, gb_pitch_, gb_roll_;
};

class CSRTParser {
  typedef enum {
    BLANK_LINE,
    SEQUENCE_NUM, 
    RELATIVE_TIME, 
    FONT_SIZE, 
    DATE_TIME,
    DETAIL_INFO,
  } PARSER_STATE_MACHINE;

public:
  CSRTParser(const std::string& fpath) {
    std::ifstream file(fpath);

    if (!file) { return ; }

    PARSER_STATE_MACHINE state = SEQUENCE_NUM;
    std::string line;
    uint32_t seq = 0;
    std::shared_ptr<FrameInfo> ptr = nullptr;
    while (std::getline(file, line)) {
      switch (state) {
      case BLANK_LINE:
        state = SEQUENCE_NUM;
        break;

      case SEQUENCE_NUM:
      {
        ptr = std::make_shared<FrameInfo>();
        seq = std::stoi(line);
        num2info_.insert(std::make_pair(seq, ptr));
        state = RELATIVE_TIME;
        break;
      }
      case RELATIVE_TIME:
      {
        int hour = 0, minute = 0, second = 0, millsec = 0;
        sscanf(line.c_str(), "%02d:%02d:%02d,%03d", &hour, &minute, &second, &millsec);
        if (ptr) {
          ptr->rel_timestamp_ = (hour * 3600 + minute * 60 + second) * 1000 + millsec;
        }
        state = FONT_SIZE;
        break;
      }
      case FONT_SIZE:
      {
        if (ptr) {
          auto idx = line.find(DIFF_TIME, 0);
          if (ptr && idx != std::string::npos) {
            ptr->diff_time_ = std::stoi(line.substr(idx + std::strlen(DIFF_TIME) + 1));
          }
        }
        state = DATE_TIME;
        break;
      }
      case DATE_TIME:
      {
        if (ptr) {
          ptr->milli_tm_ = DateTimeToMilliseconds(line);
        }
        state = DETAIL_INFO;
        break;
      }
      case DETAIL_INFO:
      {
        std::map<std::string, std::string> tokens;
        StringTokenizer(line, tokens);
        if (ptr) {
          if (tokens.find(FOCAL_LEN) != tokens.end()) 
            ptr->focal_length_ = std::stod(tokens[FOCAL_LEN]);
          if (tokens.find(LATITUDE) != tokens.end()) 
            ptr->latitude_ = std::stod(tokens[LATITUDE]);
          if (tokens.find(LONGITUDE) != tokens.end()) 
            ptr->longitude_ = std::stod(tokens[LONGITUDE]);
          if (tokens.find(REL_ALTITUDE) != tokens.end()) 
            ptr->rel_altitude_ = std::stod(tokens[REL_ALTITUDE]);
          if (tokens.find(ABS_ALTITUDE) != tokens.end()) 
            ptr->abs_alt_ = std::stod(tokens[ABS_ALTITUDE]);
          if (tokens.find(YAW) != tokens.end()) 
            ptr->gb_yaw_ = std::stod(tokens[YAW]);
          if (tokens.find(PITCH) != tokens.end()) 
            ptr->gb_pitch_ = std::stod(tokens[PITCH]);
          if (tokens.find(ROLL) != tokens.end()) 
            ptr->gb_roll_ = std::stod(tokens[ROLL]);
        }
          
        state = BLANK_LINE;
        break;
      }
      }
    }
    file.close();
    AINFO("{} records.", num2info_.size());
  }

  const std::shared_ptr<FrameInfo>& GetFrameInfo(uint32_t seq) {
    if (num2info_.find(seq) != num2info_.end()) {
      return num2info_[seq];
    }

    return nullptr;
  } 
private:
  std::chrono::milliseconds DateTimeToMilliseconds(const std::string& datetimeStr) {
    std::tm tm = {};
    std::istringstream ss(datetimeStr);
    
    // 解析日期时间字符串（格式：YYYY-MM-DD HH:MM:SS.sss）
    ss >> std::get_time(&tm, "%Y-%m-%d %H:%M:%S");
    if (ss.fail()) {
      AERROR("Failed to parse datetime string");
      return std::chrono::milliseconds(0);
    }

    // 解析毫秒部分（如果有）
    double ms = 0.0;
    if (ss.peek() == '.') {
      ss >> ms; // 读取 .sss
    }

    // 将 std::tm 转换为 time_t（UTC 时间）
    std::time_t time = std::mktime(&tm);
    if (time == -1) {
      AERROR("Invalid tm");
      return std::chrono::milliseconds(0);
    }

    // 转换为 system_clock::time_point
    auto tp = std::chrono::system_clock::from_time_t(time) + 
              std::chrono::duration_cast<std::chrono::milliseconds>(
                  std::chrono::duration<double>(ms));

    // 计算从 Unix Epoch（1970-01-01）到该时间点的毫秒数
    return std::chrono::duration_cast<std::chrono::milliseconds>(tp.time_since_epoch());
  }

  void StringTokenizer(const std::string& str, std::map<std::string, std::string>& tokens) {
    auto idx1 = str.find('[');
    auto idx2 = str.find(']');
    std::string tmp(str);

    while (std::string::npos != idx1 && std::string::npos != idx2) {
      std::string key = tmp.substr(idx1 + 1, idx2 - idx1 - 1);
      
      while (key.find("  ") != std::string::npos) {
        key.replace(key.find("  "), 2, " ");
      }
      while (key.find(" :") != std::string::npos) {
        key.replace(key.find(" :"), 2, ":");
      }
      while (key.find(": ") != std::string::npos) {
        key.replace(key.find(": "), 2, ":");
      }
      
      tmp = tmp.substr(idx2 + 1);
      std::string value;

      std::stringstream ss(key);
      while (std::getline(ss, value, ' ')) {
        std::stringstream ss(value);
        if (std::getline(ss, key, ':')) {
            if (std::getline(ss, value, ' ')) {
                tokens.insert(std::make_pair(key, value));
            }
        }
      }

      idx1 = tmp.find('[');
      idx2 = tmp.find(']');
    }
  }

private:
  const char *FRAME_CNT = "FrameCnt";
  const char *DIFF_TIME = "DiffTime";
  const char *FOCAL_LEN = "focal_len";
  const char *LATITUDE = "latitude";
  const char *LONGITUDE = "longitude";
  const char *REL_ALTITUDE = "rel_alt";
  const char *ABS_ALTITUDE = "abs_alt";
  const char *YAW = "gb_yaw";
  const char *PITCH = "gb_pitch";
  const char *ROLL = "gb_roll";
  std::unordered_map<uint32_t, std::shared_ptr<FrameInfo>> num2info_;
};

}

int main(int argc, char *argv[]) {
  if (argc < 1) {
    std::cout << "Usage: " << argv[0] << " <video file> [output path] [skip frame num]" << std::endl;
    return 0;
  }

  Exiv2FrameReader exiv2_reader_;
  std::shared_ptr<CSRTParser> parser = nullptr;
  std::string video_filepath = argv[1], srt_filepath = argv[1];
  // 
  {
    auto idx = srt_filepath.rfind('.');
    if (idx != std::string::npos) {
      std::string suffix = srt_filepath.substr(idx + 1);
      srt_filepath.replace(idx + 1, suffix.length(), "SRT");
    }
    AINFO("SRT : {}", srt_filepath);
  }
  boost::filesystem::path video_path(video_filepath);
  boost::filesystem::path srt_path(srt_filepath);
  std::string out_image_path = video_path.stem().string();
  if (argc > 2) {
    out_image_path = argv[2];
  }

  if (!fs::exists(video_path)) {
    AERROR("Video file not exists : {}", video_filepath);
    return -1;
  }

  if (!fs::exists(srt_path)) {
    AERROR("Video SRT file not exists : {}", srt_filepath);
  } else {
    parser.reset(new CSRTParser(srt_filepath));
  }

  if (!fs::exists(out_image_path)) {
    if(fs::create_directory(out_image_path)) {
    } else{
      AERROR("Error creating output image path : {}", out_image_path);
      return -1;
    }
  }
    
  cv::VideoCapture videoCapture(video_filepath);
  if (!videoCapture.isOpened()) {
    AERROR("Can't open video file: {}", video_filepath);
    return -1;
  }

  // 获取视频文件帧率
  double fps = videoCapture.get(cv::CAP_PROP_FPS);
  // 获取视频总帧数
  int total_frame_count = videoCapture.get(cv::CAP_PROP_FRAME_COUNT);
  // 每秒抽2帧
  int frame_skip_num = static_cast<int>(fps + 1) / 2. > (total_frame_count / 298.) ? 
      static_cast<int>(fps + 1) / 2 : static_cast<int>(total_frame_count / 298. + 0.5);
  if (argc > 3) {
    frame_skip_num = std::stoi(argv[3]);
  }
  AINFO("Video Framerate = {}, Skip number = {}", fps, frame_skip_num);
  int frame_number = 0;
  bool retval = true;
  cv::Mat frame;
  while (frame_number < total_frame_count) {
#if 0
    retval = videoCapture.read(frame);
    while (!retval && frame_number < total_frame_count - 1) {
      AERROR("curretn frame no = {}\tnext frame no = {}", 
            frame_number, videoCapture.get(cv::CAP_PROP_POS_FRAMES));
      videoCapture.set(cv::CAP_PROP_POS_FRAMES, frame_number + 1);
      retval = videoCapture.retrieve(frame);
      if (!retval) ++frame_number;
    }
    frame_number++;
    if (frame_number % frame_skip_num != 0) {
      continue;
    }
#else
    videoCapture.set(cv::CAP_PROP_POS_FRAMES, frame_number);
    retval = videoCapture.retrieve(frame);
    std::string skip_indicator;
    for (int index = 1; !retval && index < frame_skip_num; ++index) {
      skip_indicator += "+";
      videoCapture.set(cv::CAP_PROP_POS_FRAMES, frame_number + index);
      retval = videoCapture.retrieve(frame);
    }
    if (!skip_indicator.empty()) {
      AERROR("{}", skip_indicator);
    }
    if (!retval) {
      frame_number += frame_skip_num;
      continue;
    }
#endif
    //cv::resize(frame, frame, cv::Size{frame.cols * resize_ratio, frame.rows * resize_ratio});
    {
      constexpr int pic_standard_number_length = 6;
      std::string tempPicName = padZeros(frame_number, pic_standard_number_length);
      tempPicName = out_image_path + "/img_" + tempPicName + ".jpg";
      cv::imwrite(tempPicName, frame);
      AINFO("[{}/{}] {}", frame_number, total_frame_count, tempPicName);
      if (parser) {
        auto frameInfo = parser->GetFrameInfo(frame_number + 1);
        if (frameInfo) {
          exiv2_reader_.AddExivInfoToImage(tempPicName, frameInfo->latitude_, 
              frameInfo->longitude_, frameInfo->abs_alt_);
        }
      }
    }
    frame_number += frame_skip_num;
  }

  return 0;
}
